import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import roc_curve, auc, accuracy_score, precision_score, recall_score, f1_score

def full_output(model, X_test, y_test, threshold=0.5, is_keras=False):
    '''
    输出完整指标
    输入:
        model: 已经训练好的模型
        X_test: 测试集特征
        y_test: 测试集标签
        threshold: 阈值, 默认为0.5
        is_keras: 是否为keras模型(神经网络类), 默认为False. keras模型predict输出的就是正类概率
    输出:
        accuracy, precision, recall, f1, auc
            accuracy: 准确率
            precision: 精准率
            recall: 召回率
            f1: f1值
            auc: auc值
        绘制roc曲线
    '''
    # predict_test: 测试集的分类预测结果, 一个ndarray
    # predict_test_prob: 测试集的正类预测概率, 一个ndarray
    if is_keras:
        predict_test_prob = model.predict(X_test).flatten()
        predict_test = np.array([1 if x > threshold else 0 for x in predict_test_prob])
    else:
        predict_test = model.predict(X_test)
        if hasattr(model, "predict_proba"):
            predict_test_prob = model.predict_proba(X_test)[:, 1]
        else:
            raise ValueError("This model does not support predict_proba. Please check if 'probability=True' is set when initializing.")

    # 计算accuracy, precision, recall, f1, auc
    accuracy = accuracy_score(y_test, predict_test)
    precision = precision_score(y_test, predict_test)
    recall = recall_score(y_test, predict_test)
    f1 = f1_score(y_test, predict_test)

    fpr, tpr, thresholds = roc_curve(y_test, predict_test_prob) 
    roc_auc = auc(fpr, tpr)

    print("Accuracy for the testing set :", accuracy)
    print("Precision for the testing set :", precision)
    print("Recall for the testing set :", recall)
    print("F1 for the testing set :", f1)
    print("AUC for the testing set :", roc_auc)

    # 绘制ROC曲线
    plt.figure(dpi=150, figsize=(5,4))
    plt.plot(fpr, tpr, color='darkorange', label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], color='navy', linestyle='--')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.legend(loc="lower right")
    plt.show()

    return accuracy, precision, recall, f1, roc_auc

def cost(model, X_test, y_test, threshold=0.5, is_keras=False):
    '''
    返回经济损失. 将负类样本预测为正类会造成5的经济损失, 将正类样本预测为负类会造成1的经济损失
    '''
    # predict_test: 测试集的分类预测结果
    if is_keras:
        predict_test_prob = model.predict(X_test).flatten()
        predict_test = [1 if x > threshold else 0 for x in predict_test_prob]
    else:
        predict_test = model.predict(X_test)
    
    cost = 0
    for predict_label, true_label in zip(predict_test, y_test):
        if predict_label == 1 and true_label == 0:
            cost += 5
        elif predict_label == 0 and true_label == 1:
            cost += 1

    return cost